#include "graph.h"
#include "node.h"
#include "parser.h"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdio>
#include <getopt.h>
#include <random>
#include <string>
#include <utility>
#include <vector>

using namespace std;

void Graph::count() {
    nE = 0;
    nS = 0;
    nI = 0;
    nR = 0;

    if (dynamic == "SIS" || dynamic == "SI") {
        for (auto &node : V) {
            nS += node.curAttr["S"];
            nI += node.curAttr["I"];
        }
    }
    if (dynamic == "SIR") {
        for (auto &node : V) {
            nS += node.curAttr["S"];
            nI += node.curAttr["I"];
            nR += node.curAttr["R"];
        }
    }
    if (dynamic == "SEIR") {
        for (auto &node : V) {
            nS += node.curAttr["S"];
            nE += node.curAttr["E"];
            nI += node.curAttr["I"];
            nR += node.curAttr["R"];
        }
    }
}

size_t Graph::randomGen(size_t range) {
    random_device r;
    default_random_engine e(r());
    uniform_int_distribution<size_t> randomNumber(0, range - 1);
    return randomNumber(e);
}

void Graph::initGraph() {
    assert(network == "ER" || network == "BA" || network == "WS");
    assert(dynamic == "SI" || dynamic == "SIR" || dynamic == "SIS" ||
           dynamic == "SEIR");
    if (network == "ER")
        ER();
    else if (network == "BA")
        BA();
    else if (network == "WS")
        WS();

    if (dynamic == "SI")
        SI();
    else if (dynamic == "SIR")
        SIR();
    else if (dynamic == "SIS")
        SIS();
    else if (dynamic == "SEIR")
        SEIR();
}

void Graph::SI() {
    assert(attr.find("alpha") != attr.end());
    for (auto &item : V) {
        item.furAttr.insert(pair<string, double>("S", 1));
        item.furAttr.insert(pair<string, double>("I", 0));
    }
    for (int i = 0; i < 2; ++i) {
        size_t x = randomGen(V.size());
        V[x].furAttr["S"] = 0;
        V[x].furAttr["I"] = 1;
    }
}
void Graph::SIR() {
    assert(attr.find("alpha") != attr.end());
    assert(attr.find("beta") != attr.end());
    for (auto &item : V) {
        item.furAttr.insert(pair<string, double>("S", 1));
        item.furAttr.insert(pair<string, double>("I", 0));
        item.furAttr.insert(pair<string, double>("R", 0));
    }
    for (int i = 0; i < 2; ++i) {
        size_t x = randomGen(V.size());
        V[x].furAttr["S"] = 0;
        V[x].furAttr["I"] = 1;
        V[x].furAttr["R"] = 0;
    }
}

void Graph::SIS() {
    assert(attr.find("alpha") != attr.end());
    assert(attr.find("beta") != attr.end());
    for (auto &item : V) {
        item.furAttr.insert(pair<string, double>("S", 1));
        item.furAttr.insert(pair<string, double>("I", 0));
    }
    for (int i = 0; i < 2; ++i) {
        size_t x = randomGen(V.size());
        V[x].furAttr["S"] = 0;
        V[x].furAttr["I"] = 1;
        V[x].furAttr["I"] = 1;
    }
}

void Graph::SEIR() {
    assert(attr.find("alpha") != attr.end());
    assert(attr.find("beta") != attr.end());
    assert(attr.find("gamma") != attr.end());
    for (auto &item : V) {
        item.furAttr.insert(pair<string, double>("S", 1));
        item.furAttr.insert(pair<string, double>("E", 0));
        item.furAttr.insert(pair<string, double>("I", 0));
        item.furAttr.insert(pair<string, double>("R", 0));
    }
    for (int i = 0; i < 2; ++i) {
        size_t x = randomGen(V.size());
        V[x].furAttr["S"] = 0;
        V[x].furAttr["E"] = 0;
        V[x].furAttr["I"] = 1;
        V[x].furAttr["R"] = 0;
    }
}

void Graph::ER() {
    assert(attr.find("node") != attr.end());
    assert(attr.find("edge") != attr.end());
    auto n = static_cast<size_t>(attr["node"]);
    V.resize(n);
    for (size_t i = 0; i < n; ++i)
        V[i].ID = i;
    E.resize(n);
    for (auto &item : E)
        item.resize(n);
    for (size_t i = 0; i < n; ++i)
        for (size_t j = 0; j < n; ++j)
            E[i][j].ID = make_pair(i, j);
    size_t count = 0;
    size_t m = static_cast<size_t>(attr["edge"]);
    while (count < m) {
        auto n1 = randomGen(n);
        auto n2 = randomGen(n);
        if (E[n1][n2].isConnected || n1 == n2)
            continue;
        ++count;
        E[n1][n2].isConnected = true;
        E[n2][n1].isConnected = true;
    }
}

void Graph::BA() {
    assert(attr.find("node") != attr.end());
    assert(attr.find("m") != attr.end());
    assert(attr.find("m0") != attr.end());
    auto n = static_cast<size_t>(attr["node"]);
    auto m = static_cast<size_t>(attr["m"]);
    auto m0 = static_cast<size_t>(attr["m0"]);
    assert(m <= m0);

    V.resize(n);
    for (size_t i = 0; i < n; ++i)
        V[i].ID = i;
    E.resize(n);
    for (auto &item : E)
        item.resize(n);
    for (size_t i = 0; i < n; ++i)
        for (size_t j = 0; j < n; ++j)
            E[i][j].ID = make_pair(i, j);

    for (size_t i = 0; i < m0 - 1; ++i) {
        E[i][i + 1].isConnected = true;
        E[i + 1][i].isConnected = true;
    }

    vector<int> degree(n);
    vector<double> p(n);
    vector<double> f(n);
    double sum = 0;
    size_t count = 0;
    for (size_t i = m0; i < n; ++i) {
        sum = 0;
        count = 0;
        for (size_t j = 0; j < n; ++j) {
            degree[j] = Degree(V[j]);
            sum += degree[j];
        }
        for (size_t j = 0; j < n; ++j)
            p[j] = degree[j] / sum;
        f = p;
        for (size_t j = 1; j < n; ++j)
            f[j] += f[j - 1];

        while (count < m) {
            double probability = random();
            auto x = upper_bound(f.begin(), f.end(), probability);
            auto temp = static_cast<size_t>(x - f.begin());
            if (E[i][temp].isConnected)
                continue;
            ++count;
            E[i][temp].isConnected = true;
            E[temp][i].isConnected = true;
        }
    }
}

void Graph::WS() {
    assert(attr.find("node") != attr.end());
    assert(attr.find("k") != attr.end());
    assert(attr.find("p") != attr.end());
    auto N = static_cast<size_t>(attr["node"]);
    auto k = static_cast<size_t>(attr["k"]);
    auto p = attr["p"];
    assert(k % 2 == 0);
    assert(N >= k);
    // 不知道这一步是干嘛的，但是现在要更改实验参数了，反正WS网络不是很重要
    // assert(static_cast<double>(k) > log(static_cast<double>(N)));
    assert(N > 2);

    V.resize(N);
    for (size_t i = 0; i < N; ++i)
        V[i].ID = i;
    E.resize(N);
    for (auto &item : E)
        item.resize(N);
    for (size_t i = 0; i < N; ++i)
        for (size_t j = 0; j < N; ++j)
            E[i][j].ID = make_pair(i, j);

    int tempk = static_cast<int>(k);
    int tempN = static_cast<int>(N);
    for (int i = 0; i < tempN; ++i) {
        for (int j = -tempk / 2; j <= -1; ++j) {
            int x = (i + j + tempN) % tempN;
            E[static_cast<size_t>(i)][static_cast<size_t>(x)].isConnected =
                true;
            E[static_cast<size_t>(x)][static_cast<size_t>(i)].isConnected =
                true;
        }
        for (int j = 1; j <= tempk / 2; ++j) {
            int x = (i + j + tempN) % tempN;
            E[static_cast<size_t>(i)][static_cast<size_t>(x)].isConnected =
                true;
            E[static_cast<size_t>(x)][static_cast<size_t>(i)].isConnected =
                true;
        }
    }

    for (size_t i = 0; i < N; ++i)
        for (size_t j = 0; j < N; ++j)
            if (E[i][j].isConnected && choice(p)) {
                E[i][j].isConnected = false;
                E[j][i].isConnected = false;
                while (true) {
                    size_t m = randomGen(N);
                    if (E[i][m].isConnected || i == m)
                        continue;
                    E[i][m].isConnected = true;
                    E[m][i].isConnected = true;
                    break;
                }
            }
}

void Graph::update() {
    ++date;
    for (auto &item : V) {
        if (dynamic == "SI")
            modelSI(item);
        if (dynamic == "SIR")
            modelSIR(item);
        if (dynamic == "SIS")
            modelSIS(item);
        if (dynamic == "SEIR")
            modelSEIR(item);
    }
    for (auto &item : V)
        item.curAttr = item.furAttr;
}

void Graph::modelSI(Node &node) {
    if (node.curAttr["I"] > 0.999999)
        for (const auto &item : E[node.ID]) {
            if (item.isConnected && V[item.ID.second].curAttr["S"] > 0.999999 &&
                choice(attr["alpha"])) {
                V[item.ID.second].furAttr["S"] = 0;
                V[item.ID.second].furAttr["I"] = 1;
            }
        }
}

void Graph::modelSIR(Node &node) {
    if (node.curAttr["I"] > 0.999999) {
        for (const auto &item : E[node.ID]) {
            if (item.isConnected && V[item.ID.second].curAttr["S"] > 0.999999 &&
                choice(attr["alpha"])) {
                V[item.ID.second].furAttr["S"] = 0;
                V[item.ID.second].furAttr["I"] = 1;
                V[item.ID.second].furAttr["R"] = 0;
            }
        }

        if (choice(attr["beta"])) {
            node.furAttr["S"] = 0;
            node.furAttr["I"] = 0;
            node.furAttr["R"] = 1;
        }
    }
}

void Graph::modelSIS(Node &node) {
    if (node.curAttr["I"] > 0.999999) {
        for (const auto &item : E[node.ID]) {
            if (item.isConnected && V[item.ID.second].curAttr["S"] > 0.999999 &&
                choice(attr["alpha"])) {
                V[item.ID.second].furAttr["S"] = 0;
                V[item.ID.second].furAttr["I"] = 1;
            }
        }

        if (choice(attr["beta"])) {
            node.furAttr["S"] = 1;
            node.furAttr["I"] = 0;
        }
    }
}

void Graph::modelSEIR(Node &node) {
    if (node.curAttr["I"] > 0.999999) {
        for (const auto &item : E[node.ID]) {
            if (item.isConnected && V[item.ID.second].curAttr["S"] > 0.999999 &&
                choice(attr["alpha"])) {
                V[item.ID.second].furAttr["S"] = 0;
                V[item.ID.second].furAttr["E"] = 1;
                V[item.ID.second].furAttr["I"] = 0;
                V[item.ID.second].furAttr["R"] = 0;
            }
        }

        if (choice(attr["gamma"])) {
            node.furAttr["S"] = 0;
            node.furAttr["E"] = 0;
            node.furAttr["I"] = 0;
            node.furAttr["R"] = 1;
        }
    }

    if (node.curAttr["E"] > 0.999999) {
        if (choice(attr["beta"])) {
            node.furAttr["S"] = 0;
            node.furAttr["E"] = 0;
            node.furAttr["I"] = 1;
            node.furAttr["R"] = 0;
        }
    }
}

bool Graph::choice(double p) {
    random_device r;
    default_random_engine e(r());
    uniform_real_distribution<double> randomNumber(0.0, 1.0);
    auto temp = randomNumber(e);
    return temp < p;
}

double Graph::random() {
    random_device r;
    default_random_engine e(r());
    uniform_real_distribution<double> randomNumber(0.0, 1.0);
    auto temp = randomNumber(e);
    return temp;
}

Graph GraphGenerator(int argc, char **argv) {
    Graph M;
    M = Parser::initGraphPara(M, argc, argv);
    M.initGraph();
    return M;
}

int Graph::Degree(const Node &node) {
    int d = 0;
    for (const auto &item : E[node.ID])
        if (item.isConnected)
            ++d;
    return d;
}
